"""
数据库连接和操作相关的函数
"""

import aiomysql
from aiomysql.cursors import DictCursor


# 数据库连接参数
DB_CONFIG = {
    'host': 'localhost',
    'user': 'root',
    'password': '123456',
    'db': 'finance',
    'charset': 'utf8mb4',
    'cursorclass': DictCursor   # 以字典返回数据
}


async def get_db_connection():
    """
    获取数据库连接
    
    Returns:
        connection: 数据库连接对象
    
    Raises:
        Exception: 连接失败时抛出异常
    """
    try:
        connection = await aiomysql.connect(**DB_CONFIG)
        return connection
    except Exception as e:
        print(f"数据库连接失败: {str(e)}")
        raise


async def close_connection(connection):
    """
    关闭数据库连接
    
    Args:
        connection: 数据库连接对象
    """
    if connection:
        connection.close()
